from server import app
from fastapi import WebSocket, WebSocketDisconnect

# WebSocket端点，需要优化根据success_queue中的task_id来查询任务信息并发送
@app.websocket("/task/ws")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    try:
        while True:
            # 添加心跳检测
            await websocket.receive_text()
            
            # 优化数据库查询逻辑
            cursor.execute("""
                SELECT * FROM tasks 
                WHERE status IN ('running','finished') 
                AND (result IS NOT NULL OR status != 'completed')
            """)
            
            updated_tasks = cursor.fetchall()
            
            for task in updated_tasks:
                await websocket.send_json({
                    "type": "update",
                    "task_id": task[1],
                    "status": task[3],
                    "result": task[4]
                })
                # 标记为已通知
                cursor.execute("UPDATE tasks SET status='completed' WHERE id=?", (task[0],))
                conn.commit()
            
            await asyncio.sleep(1)
    except WebSocketDisconnect:
        await websocket.close()
    except Exception as e:
        print(f"WebSocket error: {str(e)}")
    finally:
        conn.commit()